# -*- coding:utf8 -*-
import cv2 as cv
import numpy as np

cap = cv.VideoCapture(r"videos/vtest.avi")

# take the first frame of the video
ret, frame = cap.read()
# cv.imshow("First Frame", frame)

# setup initial location of window
c, r, w, h = 400, 250, 125, 90
track_window = (c, r, w, h)

# setup the ROI(Region Of Interest) for tracking
roi = frame[r:r+h, c:c+w]
hsv_roi = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
mask = cv.inRange(hsv_roi, np.array((0, 60, 32)), np.array((180, 255, 255)))
roi_hist = cv.calcHist([hsv_roi], [0], mask, [180], [0, 180])
cv.normalize(roi_hist, roi_hist, 0, 255, cv.NORM_MINMAX)

# setup the termination criteria, either 10 iteration or move by at least 1 pt
term_crit = (cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 1)

count = 0
while(1):
    ret, frame = cap.read()

    if ret is True:
        hsv = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
        dst = cv.calcBackProject([hsv], [0], roi_hist, [0, 180], 1)

        # apply meanshift to get the new location
        ret, track_window = cv.CamShift(dst, track_window, term_crit)

        # draw it on image
        x, y, w, h = track_window
        img2 = cv.rectangle(frame, (x, y), (x + w, y + h), 255, 2)

        count = count + 1
        cv.imshow("img2", img2)
        cv.imwrite("out/{}.jpg".format(count), img2)
    else:
        break

cv.waitKey()
cv.destroyAllWindows()
